(ns thi.ng.geom.physics.core
  #?(:cljs (:require-macros [thi.ng.math.macros :as mm]))
  (:require
   [thi.ng.math.core :as m]
   [thi.ng.geom.core :as g]
   [thi.ng.geom.vector :as v :refer [vec2 vec3]]
   [thi.ng.dstruct.core :as d]
   #?(:clj [thi.ng.math.macros :as mm])))

;; Protocols

(defprotocol ITimeStep
  (timestep [_ delta]))

(defprotocol IBehavioral
  (add-behaviors [_ bmap])
  (remove-behavior [_ id])
  (clear-behaviors [_])
  (apply-behaviors [_ delta]))

(defprotocol IConstrained
  (add-constraints [_ cmap])
  (remove-constraint [_ id])
  (clear-constraints [_])
  (apply-constraints [_ delta]))

(defprotocol IParticle
  (add-force [_ f])
  (add-velocity [_ v])
  (apply-force [_ delta])
  (clear-force [_])
  (clear-velocity [_])
  (scale-velocity [_ s])
  (position [_])
  (set-position [_ p])
  (velocity [_])
  (lock [_])
  (unlock [_])
  (locked? [_]))

(defprotocol IPhysics
  (update-particles [_ delta])
  (update-springs [_ delta]))

(defn apply-to-particle
  [p delta fns]
  (when-not (locked? p)
    (loop [f (seq fns)] (when f ((first f) p delta) (recur (next f))))))

(defn apply-to-particles
  [particles delta fns]
  (loop [ps (seq particles)]
    (when ps
      (apply-to-particle (first ps) delta fns)
      (recur (next ps)))))

;; Type implementations

(deftype VerletParticle
    #?@(:clj
        [[^:unsynchronized-mutable pos
          ^:unsynchronized-mutable prev
          ^:unsynchronized-mutable force
          ^:unsynchronized-mutable locked?
          ^:unsynchronized-mutable behaviors
          ^:unsynchronized-mutable constraints
          ^double inv-weight
          _meta]

         clojure.lang.IObj
         (meta [_] _meta)
         (withMeta
          [_ meta']
          (VerletParticle. pos prev force locked? behaviors constraints inv-weight meta'))]

        :cljs
        [[^:mutable pos
          ^:mutable prev
          ^:mutable force
          ^:mutable locked?
          ^:mutable behaviors
          ^:mutable constraints
          inv-weight
          _meta]

         IMeta
         (-meta [_] _meta)

         IWithMeta
         (-with-meta
          [_ meta']
          (VerletParticle. pos prev force locked? behaviors constraints inv-weight meta'))])

  IBehavioral
  (add-behaviors
    [_ bmap] (set! behaviors (merge behaviors bmap)) _)
  (apply-behaviors
    [_ delta] (when (seq behaviors) (apply-to-particle _ delta (vals behaviors))) _)
  (remove-behavior
    [_ id] (set! behaviors (dissoc behaviors id)) _)
  (clear-behaviors
    [_] (set! behaviors {}) _)

  IConstrained
  (add-constraints
    [_ cmap] (set! constraints (merge constraints cmap)) _)
  (apply-constraints
    [_ delta] (when (seq constraints) (apply-to-particle _ delta (vals constraints))) _)
  (remove-constraint
    [_ id] (set! constraints (dissoc constraints id)) _)
  (clear-constraints
    [_] (set! constraints {}) _)

  IParticle
  (lock
    [_] (set! locked? true) _)
  (unlock
    [_] (set! locked? false) _)
  (locked?
    [_] locked?)
  (position
    [_] pos)
  (set-position
    [_ p] (set! pos p) _)
  (velocity
    [_] (m/- pos prev))
  (add-force
    [_ f] (m/+! force f) _)
  (clear-force
    [_] (g/clear! force) _)
  (apply-force
    [_ delta]
    (let [pos' (m/madd! force (* inv-weight (* delta delta)) (m/msub pos 2.0 prev))]
      (set! prev pos)
      (set! pos pos')
      (set! force (g/clear* force)))
    _)
  (scale-velocity
    [_ s] (set! prev (m/mix prev pos s)) _)

  ITimeStep
  (timestep
    [_ delta]
    (if-not locked?
      (-> (apply-behaviors _ delta)
          (apply-force delta)
          (apply-constraints delta))
      _))

  Object
  (toString
    [_]
    (pr-str
     {:pos pos
      :prev prev
      :force force
      :locked? locked?
      :inv-weight inv-weight
      :behaviors behaviors
      :constraints constraints})))

;; Spring types

(defrecord Spring
    [^VerletParticle a
     ^VerletParticle b
     ^double rest-len
     ^double strength]

  ITimeStep
  (timestep
    [_ delta]
    (let [aw   (.-inv-weight ^VerletParticle a)
          bw   (.-inv-weight ^VerletParticle b)
          pa   (position a)
          pb   (position b)
          diff (m/- pb pa)
          dist (+ (m/mag diff) 1e-6)
          nd   (* (/ (- dist rest-len) (* dist (+ aw bw))) (* strength delta))]
      (if-not (locked? a)
        (set-position a (m/madd diff (* nd aw) pa)))
      (if-not (locked? b)
        (set-position b (m/madd! diff (* (- nd) bw) pb))))))

(defrecord PullbackSpring
    [^VerletParticle a
     ^VerletParticle b
     ^double rest-len
     ^double min-len
     ^double strength]

  ITimeStep
  (timestep
    [_ delta]
    (let [pa   (position a)
          pb   (position b)
          diff (m/- pb pa)
          dist (+ (m/mag diff) 1e-6)]
      (if (> dist min-len)
        (let [aw (.-inv-weight ^VerletParticle a)
              bw (.-inv-weight ^VerletParticle b)
              nd (* (/ (- dist rest-len) (* dist (+ aw bw))) (* strength delta))]
          (if-not (locked? a)
            (set-position a (m/madd diff (* nd aw) pa)))
          (if-not (locked? b)
            (set-position b (m/madd! diff (* (- nd) bw) pb))))))))

;; VerletPhysics sim

(defrecord VerletPhysics
    [particles springs behaviors constraints listeners drag]

  IBehavioral
  (add-behaviors
    [_ bmap] (update-in _ [:behaviors] merge bmap))
  (apply-behaviors
    [_ delta] (apply-to-particles particles delta (vals behaviors)) _)
  (remove-behavior
    [_ id] (update-in _ [:behaviors] dissoc id))
  (clear-behaviors
    [_] (assoc _ behaviors {}))

  IConstrained
  (add-constraints
    [_ cmap] (update-in _ [:constraints] merge cmap))
  (apply-constraints
    [_ delta] (apply-to-particles particles delta (vals constraints)) _)
  (remove-constraint
    [_ id] (update-in _ [:constraints] dissoc id))
  (clear-constraints
    [_] (assoc _ constraints {}))

  IPhysics
  (update-particles
    [_ delta]
    (let [drag' (* delta drag)]
      (loop [ps (seq particles)]
        (when ps
          (-> ps first (scale-velocity drag') (timestep delta))
          (recur (next ps)))))
    _)
  (update-springs
    [_ delta]
    (loop [s (seq springs)]
      (when s (timestep (first s) delta) (recur (next s))))
    _)

  ITimeStep
  (timestep
    [_ iter]
    (let [delta (/ 1.0 iter)]
      (if-let [l (get listeners :timestep)] (l _))
      (loop [i iter]
        (when (pos? i)
          (if-let [l (get listeners :iter)] (l _ i))
          (-> (apply-behaviors _ delta)
              (update-springs delta)
              (update-particles delta)
              (apply-constraints delta))
          (recur (dec i)))))
    _))

;; Constructors

(defn physics
  [{:keys [particles springs behaviors constraints drag listeners]
    :or   {particles [], springs [], behaviors {}, constraints {}
           listeners {}, drag 0.0}}]
  (VerletPhysics.
   (vec particles) (vec springs) behaviors constraints listeners drag))

(defn particle
  ([pos]
   (particle pos 1.0 false))
  ([pos weight]
   (particle pos weight false))
  ([pos weight lock?]
   (VerletParticle. pos (g/clear* pos) (g/clear* pos) lock? nil nil (/ 1.0 weight) nil)))

(defn spring
  [a b rlen strength]
  (Spring. a b (double rlen) (double strength)))

;; Behaviors

(defn gravity
  [force]
  (fn [p delta]
    (add-force p (m/* force delta))))

(defn attract
  [pos r strength]
  (let [rsq (* r r)]
    (fn [p delta]
      (let [d (m/- pos (position p))
            l (+ (m/mag-squared d) 1e-6)]
        (if (< l rsq)
          (add-force p (m/* d (/ (* (- 1.0 (/ l rsq)) (* strength delta))
                                 (Math/sqrt l)))))))))

(defn attract-particle
  [p r strength]
  (let [rsq (* r r)]
    (fn [q delta]
      (if-not (= p q)
        (let [d (m/- (position p) (position q))
              l (+ (m/mag-squared d) 1e-6)]
          (if (< l rsq)
            (add-force q (m/* d (/ (* (- 1.0 (/ l rsq)) (* strength delta))
                                   (Math/sqrt l))))))))))

(defn align
  [vel strength]
  (fn [p delta]
    (add-force p (m/subm vel (velocity p) (* strength delta)))))

;; Constraints

(defn shape-constraint*
  [pred shape]
  (fn [p _]
    (let [pos (position p)]
      (if (pred pos) (set-position p (g/closest-point shape pos))))))

(defn shape-constraint-inside
  [shape] (shape-constraint* #(not (g/contains-point? shape %)) shape))

(defn shape-constraint-outside
  [shape] (shape-constraint* #(g/contains-point? shape %) shape))

(defn shape-constraint-boundary
  [shape] (shape-constraint* identity shape))

(defn distance-constraint
  [pred p r]
  (let [rsq (* r r)]
    (fn [q _]
      (let [pos (position q)]
        (if (pred (g/dist-squared p pos) rsq)
          (set-position q (m/+! (m/normalize (m/- pos p) r) p)))))))
